#!/bin/bash
# NPU log collect script.
# Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.

CURRENT_DIR=$(pwd)
TOP_DIR=$(pwd)

ASCEND_INSTALL_PATH="/usr/local/Ascend"

os_name=$(cat /etc/*release | grep -w 'NAME' | awk -F "=" '{ print $2 }')
dev_list=$(npu-smi info -l | grep "NPU ID" | awk -F ":" '{ print $2 }')
dev_id=""

usage() {
    echo "Usage: $(basename $0)"
    echo "  or   $(basename $0) [OPTION]"
    echo "Example:"
    echo "$(basename $0)                    Collect the default content."
    echo "$(basename $0) -p /home/tmp       Specify a path to store the generated content."
    echo "$(basename $0) -off host_info     Disable host os info."
    echo "Options:"
    echo "    -h, --help        Help."
    echo "    -p, --path        Specify a path to store the generated content. If the path is inexistent, this tool will create the directory."
    echo "    --off             Disable collecting specified content."
    echo "                      <host_info>: disable collecting host os info, including environment info, pcie config, proc info, and so on."
    echo "                      <host_log>: disable collecting host log, including dmesg, messages, kern.log, syslog."
    echo "                      <device_log>: disable collecting device dump log."
    echo "                      <driver_info>: disable collecting driver info."
    echo "                      <install_info>: disable collecting ascend driver install info."
    echo "                      <network_info>: disable collecting ascend driver network configuration."
}

function parse_args()
{
    while [ "$1" != "" ]; do
        case "$1" in
            -p | --path )
                if [ -z "$2" ]; then
                    usage
                    exit 1
                elif [ ! -d "$2" ] && [ ! -h "$2" ]; then
                    echo "Warning: The specified path is inexistent, this tool will create the directory." 
                    mkdir -p "$2"
                    chmod 440 "$2"
                fi
                cd "$2"
                TOP_DIR=$(pwd)
                shift
                ;;
            --off )
                if [ -z "$2" ]; then
                    usage
                    exit 1
                fi
                shift
                ;;
            -h | --help )
                usage
                exit
                ;;
            * )
                usage
                exit 1
        esac
        shift
    done
}

function create_log_collect_dir()
{
    LOG_COLLECT_DIR=$TOP_DIR/npu_log_collect
    if [ ! -d "$LOG_COLLECT_DIR" ]; then
        mkdir -p $LOG_COLLECT_DIR
        if [ $? -ne 0 ]; then
            echo "Error: Create log collect directory failed. Please check." 
            exit 1
        fi
    fi
    LOG_COLLECT_RUNING_INFO=$LOG_COLLECT_DIR/collect_scripts_running_log
}

function show_info()
{
    echo
    echo "You are now running "$(basename $0)". This is a tool which will collect necessery information"
    echo "about your system, Linux kernel, ascend kernel module, nputools and MCU. This will help our support"
    echo "team diagnose the issue you are facing while using the product. The infomation collection process may"
    echo "take several minutes and it will generate the tar file named with timestamp under the assigned directory."
    echo "Please download and send the file to the support team, thanks."
    echo
}

function get_install_variable()
{
    INSTALL_PATH=$(cat /lib/davinci.conf | awk -F '=' '{print $2}' | head -1)
}

function collect_ascend_bug_log()
{
    echo "$(date) start collect ascend_bug_report log" | tee -a $LOG_COLLECT_RUNING_INFO
    chmod 440 $LOG_COLLECT_RUNING_INFO
    $INSTALL_PATH/driver/tools/ascend_bug_report.sh "$@" 1>> $LOG_COLLECT_RUNING_INFO 2>&1
    mv $TOP_DIR/ascend_log_* $LOG_COLLECT_DIR/
    chmod 440 $LOG_COLLECT_DIR/ascend_log_*
    echo "$(date) end collect ascend_bug_report log" | tee -a $LOG_COLLECT_RUNING_INFO
}

function collect_nputools_run_log()
{
    echo "$(date) start collect run log" | tee -a $LOG_COLLECT_RUNING_INFO
    mkdir -p $LOG_COLLECT_DIR/nputools_log/root
    cd $LOG_COLLECT_DIR/nputools_log/root
    cp /var/log/nputools_LOG_ERR.log ./ 1>> $LOG_COLLECT_RUNING_INFO 2>&1
    cp /var/log/nputools_LOG_ERR.log.1 ./ 1>> $LOG_COLLECT_RUNING_INFO 2>&1
    cp /var/log/nputools_LOG_ERR.log-[0-9][0-9][0-9][0-9][0-9][0-9][0-9][0-9] ./ 1>> $LOG_COLLECT_RUNING_INFO 2>&1
    cp /var/log/nputools_LOG_ERR.log.bak ./ 1>> $LOG_COLLECT_RUNING_INFO 2>&1
    cp /var/log/nputools_LOG_INFO.log ./ 1>> $LOG_COLLECT_RUNING_INFO 2>&1
    cp /var/log/nputools_LOG_INFO.log.1 ./ 1>> $LOG_COLLECT_RUNING_INFO 2>&1
    cp /var/log/nputools_LOG_INFO.log-[0-9][0-9][0-9][0-9][0-9][0-9][0-9][0-9] ./ 1>> $LOG_COLLECT_RUNING_INFO 2>&1
    cp /var/log/nputools_LOG_INFO.log.bak ./ 1>> $LOG_COLLECT_RUNING_INFO 2>&1
    cp /var/log/nputools_LOG_OP.log ./ 1>> $LOG_COLLECT_RUNING_INFO 2>&1
    cp /var/log/nputools_LOG_OP.log.1 ./ 1>> $LOG_COLLECT_RUNING_INFO 2>&1
    cp /var/log/nputools_LOG_OP.log-[0-9][0-9][0-9][0-9][0-9][0-9][0-9][0-9] ./ 1>> $LOG_COLLECT_RUNING_INFO 2>&1
    cp /var/log/nputools_LOG_OP.log.bak ./ 1>> $LOG_COLLECT_RUNING_INFO 2>&1
    cd $CURRENT_DIR
    user_name_list=$(dir /home)
    for user_name in $user_name_list
    do
        if [ ! -d "/home/$user_name/var/log" ]
        then
            continue
        else
            mkdir -p $LOG_COLLECT_DIR/nputools_log/$user_name
            cd $LOG_COLLECT_DIR/nputools_log/$user_name
            cp /home/$user_name/var/log/nputools_LOG_ERR.log ./ 1>> $LOG_COLLECT_RUNING_INFO 2>&1
            cp /home/$user_name/var/log/nputools_LOG_ERR.log.bak ./ 1>> $LOG_COLLECT_RUNING_INFO 2>&1
            cp /home/$user_name/var/log/nputools_LOG_INFO.log ./ 1>> $LOG_COLLECT_RUNING_INFO 2>&1
            cp /home/$user_name/var/log/nputools_LOG_INFO.log.bak ./ 1>> $LOG_COLLECT_RUNING_INFO 2>&1
            cp /home/$user_name/var/log/nputools_LOG_OP.log ./ 1>> $LOG_COLLECT_RUNING_INFO 2>&1
            cp /home/$user_name/var/log/nputools_LOG_OP.log.bak ./ 1>> $LOG_COLLECT_RUNING_INFO 2>&1
            cd $CURRENT_DIR
        fi
    done
    echo "$(date) end collect run log" | tee -a $LOG_COLLECT_RUNING_INFO
}

function collect_npu_info_log()
{
    echo "$(date) start collect npu info log" | tee -a $LOG_COLLECT_RUNING_INFO

    local chip_list=$(npu-smi info -l | grep "Chip Count" | awk -F ":" '{ print $2 }')
    local card_cnt=$(npu-smi info -l | grep -E "Card|Total Count" | awk -F ":" '{ print $2 }')
    mkdir $LOG_COLLECT_DIR/npu_info_log
    touch $LOG_COLLECT_DIR/npu_info_log/card_info.log && chmod 440 $LOG_COLLECT_DIR/npu_info_log/card_info.log
    npu-smi info 1>> $LOG_COLLECT_DIR/npu_info_log/card_info.log 2>> $LOG_COLLECT_RUNING_INFO
    npu-smi info -l 1>> $LOG_COLLECT_DIR/npu_info_log/card_info.log 2>> $LOG_COLLECT_RUNING_INFO

    OLD_IFS="$IFS"
    IFS=" "
    chip_count_list=($chip_list)
    IFS="$OLD_IFS"
    j=0
    for dev_id in $dev_list
    do
        echo "########### collect npu id $dev_id product info ###########" >> $LOG_COLLECT_DIR/npu_info_log/card_info.log 
        echo "npu-smi info -t product -i $dev_id -c 0 1" >> $LOG_COLLECT_DIR/npu_info_log/card_info.log
        npu-smi info -t product -i $dev_id -c 0 1>> $LOG_COLLECT_DIR/npu_info_log/card_info.log 2>> $LOG_COLLECT_RUNING_INFO

        for ((i=0;i<${chip_count_list[j]};i++))
        do
            if [ $j -ge 0 ] && [ $j -lt $card_cnt ]; then
                echo "############ collect npu id $dev_id chip id $i board info ############" >> $LOG_COLLECT_DIR/npu_info_log/card_info.log
                echo "npu-smi info -t board -i $dev_id -c $i" >> $LOG_COLLECT_DIR/npu_info_log/card_info.log
                npu-smi info -t board -i $dev_id -c $i 1>> $LOG_COLLECT_DIR/npu_info_log/card_info.log 2>> $LOG_COLLECT_RUNING_INFO

                echo "########### collect npu id $dev_id chip id $i health info #############" >> $LOG_COLLECT_DIR/npu_info_log/card_info.log
                echo "npu-smi info -t health -i $dev_id -c $i" >> $LOG_COLLECT_DIR/npu_info_log/card_info.log
                npu-smi info -t health -i $dev_id -c $i 1>> $LOG_COLLECT_DIR/npu_info_log/card_info.log 2>> $LOG_COLLECT_RUNING_INFO
            else
                break
            fi
        done
        ((j++))
    done
    echo "$(date) end collect npu info log" | tee -a $LOG_COLLECT_RUNING_INFO
}

function collect_mcu_log()
{
    echo "$(date) start collect mcu_log" | tee -a $LOG_COLLECT_RUNING_INFO

    if [ -d "/run/mcu_log" ]; then
        rm -f /run/mcu_log/error_log*.log /run/mcu_log/maintaince_log*.log /run/mcu_log/operate_log*.log 1>> $LOG_COLLECT_RUNING_INFO 2>&1
    fi

    for dev_id in $dev_list
    do
        npu-smi set -t collect-log -i $dev_id 1>> $LOG_COLLECT_RUNING_INFO 2>&1 &
    done
    wait

    for dev_id in $dev_list
    do
        mkdir -p $LOG_COLLECT_DIR/mcu_log/card_${dev_id}_mcu_log
        cp -f /run/mcu_log/error_log_${dev_id}_*.log /run/mcu_log/maintaince_log_${dev_id}_*.log /run/mcu_log/operate_log_${dev_id}_*.log $LOG_COLLECT_DIR/mcu_log/card_${dev_id}_mcu_log/ 1>> $LOG_COLLECT_RUNING_INFO 2>&1
    done
    echo "$(date) end collect mcu_log" | tee -a $LOG_COLLECT_RUNING_INFO
}

function compress_file()
{
    START_TIME=$(date +"%Y-%m-%d %H:%M:%S")
    DIR_DATE=$(echo $START_TIME | tr -cd "[0-9]")
    cd $TOP_DIR
    echo "compressing files into a tar package" | tee -a $LOG_COLLECT_RUNING_INFO
    if [ "$os_name" == "SLES" ];then
        tar  --format=gnu -czPf npu_log_collect_$DIR_DATE.tar.gz npu_log_collect
    else
        tar -czPf npu_log_collect_$DIR_DATE.tar.gz npu_log_collect
    fi
    chmod 440 npu_log_collect_$DIR_DATE.tar.gz
    rm -rf $LOG_COLLECT_DIR
    cd
}

pid=$$
function conflict_check()
{
    PS_INFO=$(ps -ef)
    getpid=$(echo "$PS_INFO" | grep "npu_log_collect.sh" | grep -v "grep" | awk '{print $2}')
    for loop in $getpid
    do
        if [ "$loop" = "$pid" ]; then
            exitpid=0
        else
            exitpid=1
            break
        fi
    done

    if [ $exitpid -eq 1 ]; then
        echo "Error: Other npu_log_collect.sh task is running."
        exit 1
    fi
}

function main()
{
    show_info
    get_install_variable
    create_log_collect_dir
    collect_ascend_bug_log "$@"
    collect_nputools_run_log
    collect_npu_info_log
    collect_mcu_log
    compress_file
}

function check_run_in_vm()
{
    dmidecode | grep -E "Manufacturer: QEMU|Manufacturer: qemu|xen|Xen|VMware|OpenStack|KVM"
    if [ $? -eq 0 ]; then
        return 1
    else
        /usr/bin/systemd-detect-virt > /dev/null  2>&1
        if [ $? -eq 0 ]; then
            return 1
        else
            return 0
        fi
    fi
}

function check_run_in_docker()
{
    cat /proc/self/cgroup | grep "docker"
    if [ $? -eq 0 ]; then
        return 1
    else
        return 0
    fi
}


if [ $(id -u) -ne 0 ]; then
    echo "Error: Please run $(basename $0) as root."
    exit 1
fi

check_run_in_vm
if [ $? -eq 1 ]; then
    echo "Error: Please run $(basename $0) on physical machine."
    exit 1
fi

check_run_in_docker
if [ $? -eq 1 ]; then
    echo "Error: Please run $(basename $0) on physical machine."
    exit 1
fi

conflict_check
parse_args "$@"
main "$@"

